import pickle
import re
import joblib
import pandas as pd
import psg
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics import classification_report
from sklearn.naive_bayes import MultinomialNB
from zhconv import zhconv


def train():
    train_data = pickle.load(open("data.03-模型训练数据.pkl","rb"))
    model =MultinomialNB()
    model.fit(train_data["x"],train_data["y"])

    joblib.dump(model,"04-垃圾邮件分类模型.pth")

def clean_data(content):
    # 2.取出中文
    content = re.sub(r"[\u4e00-\u9fa5]"," ",content)
    # 3.繁体转为简体
    content = zhconv.convert(content,"zh-cn")
    # 4.分词. 过滤词性
    content_pos = psg.cut(content)
    # 需要保留的词性
    allow_pos = ['n', 'nr', 'ns', 'nt', 'v', 'a']
    words = []
    for word,pos in content_pos:
        if pos in allow_pos:
            words.append(word)
    # 转换成str类型
    return " ".join(content)

# 模型评估
def evaluate():
    # 1.加载原始数据集
    test_data = pd.read_csv("data/01.原始训练集.csv")

    vocab = pickle.load(open("data/03-模型训练特征.pkl","rb"))
    transfer = CountVectorizer(vocabulary=vocab)
    x_test,y_test = [], []
    for content,label in zip(test_data["content"],test_data["label"]):
        #2.清洗数据
        content = clean_data(content)
        if len(content) == 0:
            continue
        # 3.特征功能
        x_test.append(transfer.transform(content))
        y_test.append(label)
    # 4.加载模型
    model = joblib.load("04-垃圾邮件分类模型.pth")
    # 5.预测
    y_pre = model.predict(x_test)
    # 6.评估
    print(classification_report(y_test,y_pre))



if __name__ == '__main__':
    # train()
    evaluate()





